'''
Created on Aug 18, 2009

@author: mkiyer
'''
"""
Thanks Josh Hemann for the example
"""

import operator
import collections
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
import sys
import os

from bx.intervals.intersection import Interval, IntervalTree

class Sample:    
    def __init__(self, name):
        self.name = name
        self.data = []


def read_expression_matrix(fhd):
    # first line is property header with sample names
    property_header = fhd.readline()
    sample_names = property_header.strip().split('\t')[1:]
    samples = [Sample(x) for x in sample_names]
    # set sample properties
    while True:
        thisline = fhd.readline()
        thisline = thisline.strip()
        if len(thisline) == 0:
            break
        fields = thisline.split('\t')
        prop_name = fields[0]
        for i,prop_val in enumerate(fields[1:]):
            setattr(samples[i], prop_name, prop_val)
    # ignore for now?
    data_header = fhd.readline()
    # get sample data
    intervals = []
    for thisline in fhd:
        fields = thisline.strip().split('\t')
        chrom, interval = fields[0].split(':')
        start, end = map(int, interval.split('-'))
        intervals.append((chrom, start, end))
        for i, data_val in enumerate(fields[1:]):
            samples[i].data.append(float(data_val))    
    return intervals, samples    

def group_samples_by_properties(samples, props):
    groups = collections.defaultdict(lambda: [])    
    for s in samples:
        prop_key = []
        for p in props:
            if hasattr(s, p):
                if getattr(s, p) != "None":
                    prop_key.append(getattr(s,p))
                else:
                    print 'sample %s has None for attribute %s' % (s, p)
            else:
                print 'sample has no attribute %s' % p
        prop_key = tuple(prop_key)
        #prop_key = tuple([getattr(s, p) for p in props])
        groups[prop_key].append(s)    
    return groups

def progression_compare(x, y):
    order = ['Benign', 'Benign Cell Line', 'Tumor', 'Tumor Cell Line', 'Metastatic']

    xi, yi = None, None
    for i, val in enumerate(order):        
        if x == val:
            xi = i
        if y == val:
            yi = i
    
    if xi != None and yi != None:
        return xi - yi

    if x>y:
        return 1
    elif x==y:
        return 0
    else: # x<y
        return -1

def draw_boxplot(data, labels, colors, prop_cmap, window_title="A Boxplot Example", title="Example interval", outfile="testboxplot.png"):
    fig = plt.figure(figsize=(20,8))
    fig.canvas.set_window_title(window_title)
    ax1 = fig.add_subplot(111)
    plt.subplots_adjust(left=0.075, right=0.95, top=0.9, bottom=0.25)
    
    bp = plt.boxplot(data, notch=0, sym='+', vert=1, whis=1.5)
    plt.setp(bp['boxes'], color='black')
    plt.setp(bp['whiskers'], color='black')
    plt.setp(bp['fliers'], color='red', marker='+')
    
    # Add a horizontal grid to the plot, but make it very light in color
    # so we can use it for reading data values but not be distracting
    ax1.yaxis.grid(True, linestyle='-', which='major', color='lightgrey',
                  alpha=0.5)
    
    # Hide these grid behind plot objects
    ax1.set_axisbelow(True)
    ax1.set_title(title)
    ax1.set_xlabel('Sample group')
    ax1.set_ylabel('Coverage (bases covered per kilobase per million reads)')
    
    # Now fill the boxes with desired colors
    #boxColors = ['darkkhaki','royalblue']
    boxColors = colors
    numBoxes = len(data)
    medians = range(numBoxes)
    for i in range(numBoxes):
      box = bp['boxes'][i]
      boxX = []
      boxY = []
      for j in range(5):
          boxX.append(box.get_xdata()[j])
          boxY.append(box.get_ydata()[j])
      boxCoords = zip(boxX,boxY)
      # Alternate between Dark Khaki and Royal Blue
      #k = i % 2
      boxPolygon = Polygon(boxCoords, facecolor=boxColors[i])
      ax1.add_patch(boxPolygon)
      # Now draw the median lines back over what we just filled in
      med = bp['medians'][i]
      medianX = []
      medianY = []
      for j in range(2):
          medianX.append(med.get_xdata()[j])
          medianY.append(med.get_ydata()[j])
          plt.plot(medianX, medianY, 'k')
          medians[i] = medianY[0]
      # Finally, overplot the sample averages, with horizontal alignment
      # in the center of each box
      plt.plot([np.average(med.get_xdata())], [np.average(data[i])],
               color='w', marker='*', markeredgecolor='k')

    # Set the axes ranges and axes labels
    ax1.set_xlim(0.5, numBoxes+0.5)
    # TODO: setting ylim appears to screw things up
    bottom, top = ax1.get_ylim()
    #top = 1000
    #bottom = -5
    #ax1.set_ylim(bottom, top)
    xtickNames = plt.setp(ax1, xticklabels=labels)
    plt.setp(xtickNames, rotation=45, fontsize=8)
    
    # Due to the Y-axis scale being different across samples, it can be
    # hard to compare differences in medians across the samples. Add upper
    # X-axis tick labels with the sample medians to aid in comparison
    # (just use two decimal places of precision)
    pos = np.arange(numBoxes)+1
    upperLabels = [str(np.round(s, 2)) for s in medians]
    weights = ['bold', 'semibold']
    for tick,label in zip(range(numBoxes),ax1.get_xticklabels()):
       k = tick % 2
       ax1.text(pos[tick], top-(top*0.05), upperLabels[tick],
            horizontalalignment='center', size='x-small', weight=weights[k],
            color=boxColors[tick])
    
    # Finally, add a basic legend
    ystart = 0.045
    yincrement = 0.02
    for color_key, color in prop_cmap.iteritems():
        plt.figtext(0.80, ystart, color_key,
                    backgroundcolor=color, color='black', weight='roman',
                    size='x-small')
        ystart += yincrement
    plt.figtext(0.80, 0.015, '*', color='white', backgroundcolor='silver',
               weight='roman', size='medium')
    plt.figtext(0.815, 0.013, ' Average Value', color='black', weight='roman',
               size='x-small')
    #plt.show()
    plt.savefig(outfile)
    # close or memory leak will happen
    plt.close()

def color_by_prop(group_keys, getpropfunc):    
    import matplotlib.cm as cm
    props = [getpropfunc(group_key) for group_key in group_keys]
    sorted_props = sorted(set(props), cmp=progression_compare)
    prop_colors = {}
    for i, prop in enumerate(sorted_props):
        prop_colors[prop] = cm.gist_rainbow(i * cm.gist_rainbow.N / len(sorted_props))
    return prop_colors, [prop_colors[x] for x in props]

if __name__ == '__main__':    
    # expression file
    expr_file = sys.argv[1]
    
    # get and organize data
    intervals, samples = read_expression_matrix(open(expr_file))
    group_props = ['Progression', 'TissueType', 'Treatment', 'ETS', 'ERBB2']
    groups = group_samples_by_properties(samples, group_props)
    num_groups = len(groups)
    # debug number of groups
    for k, v in groups.iteritems():
        print 'group %s -> %d' % (str(k), len(v))

    # order groups based on a key property
    getpropfunc = operator.itemgetter(0)
    # sort groups
    sorted_keys = sorted(groups.keys(), key=getpropfunc, cmp=progression_compare)
    prop_cmap, colors = color_by_prop(sorted_keys, getpropfunc)
    sorted_names = [','.join(key) for key in sorted_keys]

    # interval picture path
    dirname = os.path.splitext(os.path.basename(expr_file))[0]
    if not os.path.exists(dirname):
        print 'creating path: %s' % dirname
        os.makedirs(dirname)
    resultpath = os.path.join('.', dirname)
    
    # generate a plot for each interval
    for i, interval in enumerate(intervals):        
        data = [np.array([s.data[i] for s in groups[group_key]]) 
                for group_key in sorted_keys]
        outfile = os.path.join(resultpath, '%s_%d-%d.png' % (interval[0], interval[1], interval[2]))
        print 'drawing boxplot for %s:%d-%d' % (interval[0], interval[1], interval[2])
        draw_boxplot(data, sorted_names, colors=colors, prop_cmap=prop_cmap, window_title="Interval Boxplot", title="%s:%d-%d" % (interval[0], interval[1], interval[2]), outfile=outfile)
    


